# -!-encoding:utf-8-!-
import torch
from transformers import BertModel, BertTokenizer, BertConfig, RobertaTokenizer, RobertaModel


def t0():
    name_or_path = r"C:\Users\HP\dataroot\models\bert-base-chinese"
    tokenizer = BertTokenizer.from_pretrained(
        pretrained_model_name_or_path=name_or_path
    )
    model = BertModel.from_pretrained(
        config=BertConfig.from_pretrained(name_or_path),
        pretrained_model_name_or_path=name_or_path
    )

    texts = [
        "这个酒店的服务非常不错，curverade下次还过来",
        "我不喜欢chi这里的food食物"
    ]
    result = tokenizer(texts, padding=True)
    print(result)

    print(tokenizer._tokenize(texts[0]))
    print(tokenizer._tokenize(texts[1]))
    print(tokenizer._tokenize(
        '， 患 者 2 0 0 8 年 9 月 3 日 因 “ 腹 胀 ， 发 现 腹 部 包 块 ” 在 我 院 腹 科 行 手 术 探 查 ， 术 中 见 盆 腹 腔 肿 物 ， 与 肠 管 及 子 宫 关 系 密 切 ， 遂 行 “ 全 子 宫 左 附 件 切 除 + 盆 腔 肿 物 切 除 + 右 半 结 肠 切 除 + D I X O N 术 ” ， 术 后 病 理 示 颗 粒 细 胞 瘤 ， 诊 断 为 颗 粒 细 胞 瘤 I I I C 期 ， 术 后 自 2 0 0 8 年 1 1 月 起 行 B E P 方 案 化 疗 共 4 程 ， 末 次 化 疗 时 间 为 2 0 0 9 年 3 月 2 6 日 。 之 后 患 者 定 期 复 查 ， 2 0 1 5 - 6 - 1 ， 复 查 C T 示 ： 髂 嵴 水 平 上 腹 部 L 5 腰 椎 前 见 软 组 织 肿 块 ， 大 小 约 3 0 M M × 4 5 M M ， 密 度 欠 均 匀 ， 边 界 尚 清 楚 ， 轻 度 强 化 。 查 肿 瘤 标 志 物 均 正 常 。 于 2 0 1 5 - 7 - 6 行 剖 腹 探 查 + 膀 胱 旁 肿 物 切 除 + 骶 前 肿 物 切 除 + 肠 表 面 肿 物 切 除 术 ， 术 程 顺 利 ， ， 术 后 病 理 示 ： 膀 胱 旁 肿 物 及 骶 前 肿 物 符 合 颗 粒 细 胞 瘤 。 于 2 0 1 5 - 7 - 1 3 、 8 - 1 4 给 予 泰 素 2 4 0 M G + 伯 尔 定 6 0 0 M G 化 疗 2 程 ， 过 程 顺 利 。 出 院 至 今 ， 无 发 热 ， 无 腹 痛 、 腹 胀 ， 有 脱 发 ， 现 返 院 复 诊 ， 拟 行 再 次 化 疗 收 入 院 。 起 病 以 来 ， 精 神 、 胃 纳 、 睡 眠 可 ， 大 小 便 正 常 ， 体 重 无 明 显 改 变 。'))

    input_ids = torch.tensor(result['input_ids'], dtype=torch.long)
    attention_mask = torch.tensor(result['attention_mask'], dtype=torch.float32)

    result = model(
        input_ids=input_ids,
        attention_mask=attention_mask
    )
    print(result)


def t1():
    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, BartForConditionalGeneration, \
        AutoModelForTokenClassification, AutoModelForSequenceClassification

    name_or_path = r"C:\Users\HP\dataroot\models\bert-base-chinese"
    tokenizer = AutoTokenizer.from_pretrained(
        pretrained_model_name_or_path=name_or_path
    )
    model = AutoModelForTokenClassification.from_pretrained(
        pretrained_model_name_or_path=name_or_path
    )

    print(tokenizer)
    print(type(tokenizer))
    print(type(model))

    texts = [
        "这个酒店的服务非常不错，curverade下次还过来",
        "我不喜欢chi这里的food食物"
    ]
    result = tokenizer(texts, padding=True)
    print(result)

    input_ids = torch.tensor(result['input_ids'], dtype=torch.long)
    attention_mask = torch.tensor(result['attention_mask'], dtype=torch.float32)

    result = model(
        input_ids=input_ids,
        attention_mask=attention_mask
    )
    print(result)


def t2():
    from transformers import AlbertModel, AlbertTokenizer, AlbertConfig

    name_or_path = r"C:\Users\HP\dataroot\models\albert"

    tokenizer = AlbertTokenizer.from_pretrained(
        pretrained_model_name_or_path=name_or_path
    )
    model = AlbertModel.from_pretrained(
        pretrained_model_name_or_path=name_or_path
    )
    texts = [
        "这个酒店的服务非常不错，curverade下次还过来",
        "我不喜欢chi这里的food食物"
    ]
    result = tokenizer(texts, padding=True)
    print(result)

    input_ids = torch.tensor(result['input_ids'], dtype=torch.long)
    attention_mask = torch.tensor(result['attention_mask'], dtype=torch.float32)

    result = model(
        input_ids=input_ids,
        attention_mask=attention_mask
    )
    print(result)


def t3():
    from transformers import DebertaForTokenClassification, DebertaModel, DebertaConfig, DebertaTokenizer

    name_or_path = "microsoft/deberta-base"
    # from huggingface_hub import snapshot_download
    # snapshot_download(name_or_path) # 下载所有
    tokenizer = DebertaTokenizer.from_pretrained(
        pretrained_model_name_or_path=name_or_path
    )
    model = DebertaModel.from_pretrained(
        pretrained_model_name_or_path=name_or_path
    )
    model = DebertaModel(config=DebertaConfig())
    input_ids = torch.tensor([
        [1, 2, 3, 4, 5],
        [2, 5, 6, 9, 0]
    ], dtype=torch.long)
    attention_mask = torch.tensor([
        [1.0, 1.0, 1.0, 1.0, 1.0],
        [1.0, 1.0, 1.0, 1.0, 0.0],
    ])

    # noinspection PyTypeChecker
    outputs = model.forward(
        input_ids=input_ids,
        attention_mask=attention_mask,
        output_hidden_states=True,  # 是否返回所有层的hidden_states
        return_dict=False  # 以对象的形式返回，默认就是True
    )
    print(outputs)
    print(len(outputs))


def t4():
    name_or_path = r"C:\Users\HP\dataroot\models\bert-base-chinese"
    # tokenizer = RobertaTokenizer.from_pretrained(
    #     pretrained_model_name_or_path=name_or_path
    # )

    # 由于模型参数是bert模型参数，但是当前模型使用的是roberta
    checkpoint_file = rf"{name_or_path}\pytorch_model.bin"
    ckpt = torch.load(checkpoint_file, map_location="cpu")
    ckpt = {k[5:]: v for k, v in ckpt.items()}
    model = RobertaModel.from_pretrained(
        config=BertConfig.from_pretrained(name_or_path),
        pretrained_model_name_or_path=name_or_path,
        state_dict=ckpt
    )

    # texts = [
    #     "这个酒店的服务非常不错，curverade下次还过来",
    #     "我不喜欢chi这里的food食物"
    # ]
    # result = tokenizer(texts, padding=True)
    # print(result)
    # input_ids = torch.tensor(result['input_ids'], dtype=torch.long)
    # attention_mask = torch.tensor(result['attention_mask'], dtype=torch.float32)

    input_ids = torch.tensor([
        [1, 2, 3, 4, 5],
        [2, 5, 6, 9, 0]
    ], dtype=torch.long)
    attention_mask = torch.tensor([
        [1.0, 1.0, 1.0, 1.0, 1.0],
        [1.0, 1.0, 1.0, 1.0, 0.0],
    ])
    result = model(
        input_ids=input_ids,
        attention_mask=attention_mask
    )
    print(result)


if __name__ == '__main__':
    t0()
